import cv2
import numpy as np
import de
import os
import math
from shapely.geometry import Polygon
from .base_tpu import TPU


CLASSES = ['rect']


class PredictObb(TPU):
    def __init__(self):
        super().__init__()
        # 提取文件名, 输出对应的文件路径
        file_name = os.path.basename(__file__)
        print(f"{file_name}.py path : {self.pwd}")

    def predict(self, img, imgsz, conf=None, iou=None):
        data = self.cap_processing(img, imgsz)
        output = self.engine.predict(data)
        # results_box = []
        if conf or iou:
            results = self.outbox_processing(output, conf, iou)
            if results is not None:
                return results

    # 视频画面处理
    def cap_processing(self, img, imgsz):
        format = de.PixelFormat.DE_PIX_FMT_JPEG
        imgsz = tuple(imgsz)
        # print(f"imgsz:{imgsz}")
        letterbox_img, self.aspect_ratio, self.offset_x, self.offset_y = letterbox_resize(img, imgsz, 114)
        encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), 20]
        result, encoded_image = cv2.imencode('.jpg', letterbox_img, encode_param)
        img = np.array(encoded_image)
        shape = (1, 1, 1, len(img))
        data = [(format, shape, img)]
        return data

    def outbox_processing(self, results, conf, input_iou):
        outputs=[]
        results_list = []
        for x in results[:-1]:
            index, stride=0, 0
            if x.shape[2]==20:
                stride=32
                index=20*4*20*4+20*2*20*2
            if x.shape[2]==40:
                stride=16
                index=20*4*20*4
            if x.shape[2]==80:
                stride=8
                index=0
            feature=x.reshape(1,65,-1)
            
            output=process(feature,x.shape[3],x.shape[2],stride,results[-1],index, conf)
            outputs=outputs+output
            
        predbox = NMS(outputs, input_iou)
        
        for index in range(len(predbox)):
            xmin = int((predbox[index].xmin-self.offset_x)/self.aspect_ratio)
            ymin = int((predbox[index].ymin-self.offset_y)/self.aspect_ratio)
            xmax = int((predbox[index].xmax-self.offset_x)/self.aspect_ratio)
            ymax = int((predbox[index].ymax-self.offset_y)/self.aspect_ratio)
            classId = predbox[index].classId
            score = predbox[index].score
            angle = predbox[index].angle
            xywh_ = predbox[index].xywh_
            # print("w,h========",xywh_)
            # ptext = (xmin, ymin)
            points=rotate_rectangle(xmin,ymin,xmax,ymax,angle)
            # print("points========",points)
            w= xywh_[2] 
            h = xywh_[3] 
            # h,w = wdandle(points)
            angle = -angle
            # theta = angle / 180 * math.pi
            # print("width+++++++++",w)
            # print("height+++++++++",h)
            if w != max(w, h): 
                w, h = h, w
#                angle = -angle
                angle = angle + math.pi/2

            if angle < 0:
                angle = math.pi - abs(angle)

            angle = int(angle / math.pi * 180)
            # print("angle = ", angle)
            center_x = int((xmin + xmax) / 2)  # int(xywh_[0])
            center_y = int((ymin + ymax) / 2) # int(xywh_[1] * 480 / 640)

            results_list.append([center_x, center_y, int(w), int(h), score, classId, angle])
        return results_list


class DetectBox:
    def __init__(self, classId, score, xmin, ymin, xmax, ymax,angle,xywh_):
        self.classId = classId
        self.score = score
        self.xmin = xmin
        self.ymin = ymin
        self.xmax = xmax
        self.ymax = ymax
        self.angle=angle
        self.xywh_ = xywh_


def letterbox_resize(image, size, bg_color):
    """
    letterbox_resize the image according to the specified size
    :param image: input image, which can be a NumPy array or file path
    :param size: target size (width, height)
    :param bg_color: background filling data 
    :return: processed image
    """
    if isinstance(image, str):
        image = cv2.imread(image)

    target_width, target_height = size
    image_height, image_width, _ = image.shape

    # Calculate the adjusted image size
    aspect_ratio = min(target_width / image_width, target_height / image_height)
    new_width = int(image_width * aspect_ratio)
    new_height = int(image_height * aspect_ratio)

    # Use cv2.resize() for proportional scaling
    image = cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_AREA)

    # Create a new canvas and fill it
    result_image = np.ones((target_height, target_width, 3), dtype=np.uint8) * bg_color
    offset_x = (target_width - new_width) // 2
    offset_y = (target_height - new_height) // 2
    result_image[offset_y:offset_y + new_height, offset_x:offset_x + new_width] = image
    return result_image, aspect_ratio, offset_x, offset_y


def rotate_rectangle(x1, y1, x2, y2, a):
    # 计算中心点坐标
    cx = (x1 + x2) / 2
    cy = (y1 + y2) / 2

    # 将角度转换为弧度
    # a = math.radians(a)
    # 对每个顶点进行旋转变换
    x1_new = int((x1 - cx) * math.cos(a) - (y1 - cy) * math.sin(a) + cx)
    y1_new = int((x1 - cx) * math.sin(a) + (y1 - cy) * math.cos(a) + cy)

    x2_new = int((x2 - cx) * math.cos(a) - (y2 - cy) * math.sin(a) + cx)
    y2_new = int((x2 - cx) * math.sin(a) + (y2 - cy) * math.cos(a) + cy)

    x3_new = int((x1 - cx) * math.cos(a) - (y2 - cy) * math.sin(a) + cx)
    y3_new = int((x1 - cx) * math.sin(a) + (y2 - cy) * math.cos(a) + cy)

    x4_new =int( (x2 - cx) * math.cos(a) - (y1 - cy) * math.sin(a) + cx)
    y4_new =int( (x2 - cx) * math.sin(a) + (y1 - cy) * math.cos(a) + cy)
    return [(x1_new, y1_new), (x3_new, y3_new),(x2_new, y2_new) ,(x4_new, y4_new)]


def process(out,model_w,model_h,stride,angle_feature,index,objectThresh,scale_w=1,scale_h=1):
    class_num=len(CLASSES)
    angle_feature=angle_feature.reshape(-1)
    xywh=out[:,:64,:]
    
    conf=sigmoid(out[:,64:,:])
    out=[]
    conf=conf.reshape(-1)
    for ik in range(model_h*model_w*class_num):
        
        if conf[ik]>objectThresh:
            w=ik%model_w
            h=(ik%(model_w*model_h))//model_w
            c=ik//(model_w*model_h)
            xywh_=xywh[0,:,(h*model_w)+w] #[1,64,1]
            xywh_=xywh_.reshape(1,4,16,1)
            data=np.array([i for i in range(16)]).reshape(1,1,16,1)
            xywh_=softmax(xywh_,2)
            xywh_ = np.multiply(data, xywh_)
            xywh_ = np.sum(xywh_, axis=2, keepdims=True).reshape(-1)
            xywh_add=xywh_[:2]+xywh_[2:]
            xywh_sub=(xywh_[2:]-xywh_[:2])/2
            angle_feature_= (angle_feature[index+(h*model_w)+w]-0.25)*3.1415927410125732
            angle_feature_cos=math.cos(angle_feature_)
            angle_feature_sin=math.sin(angle_feature_)
            xy_mul1=xywh_sub[0] * angle_feature_cos
            xy_mul2=xywh_sub[1] * angle_feature_sin
            xy_mul3=xywh_sub[0] * angle_feature_sin
            xy_mul4=xywh_sub[1] * angle_feature_cos
            xy=xy_mul1-xy_mul2,xy_mul3+xy_mul4
            xywh_1=np.array([(xy_mul1-xy_mul2)+w+0.5,(xy_mul3+xy_mul4)+h+0.5,xywh_add[0],xywh_add[1]])
            xywh_=xywh_1*stride
            
            xmin = (xywh_[0] - xywh_[2] / 2) * scale_w
            ymin = (xywh_[1] - xywh_[3] / 2) * scale_h
            xmax = (xywh_[0] + xywh_[2] / 2) * scale_w
            ymax = (xywh_[1] + xywh_[3] / 2) * scale_h
            box = DetectBox(c,conf[ik], xmin, ymin, xmax, ymax,angle_feature_,xywh_)
            out.append(box)
    return out


def NMS(detectResult, nmsThresh):
    predBoxs = []

    sort_detectboxs = sorted(detectResult, key=lambda x: x.score, reverse=True)
    for i in range(len(sort_detectboxs)):
        xmin1 = sort_detectboxs[i].xmin
        ymin1 = sort_detectboxs[i].ymin
        xmax1 = sort_detectboxs[i].xmax
        ymax1 = sort_detectboxs[i].ymax
        classId = sort_detectboxs[i].classId
        angle = sort_detectboxs[i].angle
        p1=rotate_rectangle(xmin1, ymin1, xmax1, ymax1, angle)
        p1=np.array(p1).reshape(-1)
        
        if sort_detectboxs[i].classId != -1:
            predBoxs.append(sort_detectboxs[i])
            for j in range(i + 1, len(sort_detectboxs), 1):
                if classId == sort_detectboxs[j].classId:
                    xmin2 = sort_detectboxs[j].xmin
                    ymin2 = sort_detectboxs[j].ymin
                    xmax2 = sort_detectboxs[j].xmax
                    ymax2 = sort_detectboxs[j].ymax
                    angle2 = sort_detectboxs[j].angle
                    p2=rotate_rectangle(xmin2, ymin2, xmax2, ymax2, angle2)
                    p2=np.array(p2).reshape(-1)
                    iou=intersection(p1, p2)
                    if iou > nmsThresh:
                        sort_detectboxs[j].classId = -1
    return predBoxs


def sigmoid(x):
    return 1 / (1 + np.exp(-x))

def softmax(x, axis=-1):
    # 将输入向量减去最大值以提高数值稳定性
    exp_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
    return exp_x / np.sum(exp_x, axis=axis, keepdims=True)


def intersection(g, p):
    g=np.asarray(g)
    p=np.asarray(p)
    g = Polygon(g[:8].reshape((4, 2)))
    p = Polygon(p[:8].reshape((4, 2)))
    if not g.is_valid or not p.is_valid:
        return 0
    inter = Polygon(g).intersection(Polygon(p)).area
    union = g.area + p.area - inter
    if union == 0:
        return 0
    else:
        return inter/union


# 计算两点之间的距离
def distance(p1, p2):
    return math.sqrt((p2[0] - p1[0])**2 + (p2[1] - p1[1])**2)


# 计算宽（与长边垂直的方向上的最大距离）
# 假设长边是 points[0] 到 points[1]
# 计算 points[2] 和 points[3] 到长边的垂直距离
def perpendicular_distance(p, p1, p2):
    # 计算点到直线的距离公式
    numerator = abs((p2[1] - p1[1]) * p[0] - (p2[0] - p1[0]) * p[1] + p2[0] * p1[1] - p2[1] * p1[0])
    denominator = math.sqrt((p2[1] - p1[1])**2 + (p2[0] - p1[0])**2)
    return numerator / denominator


def wdandle(points):
        # 计算所有边的长度
    edges = [
        distance(points[0], points[1]),
        distance(points[1], points[2]),
        distance(points[2], points[3]),
        distance(points[3], points[0])
    ]

    # 找到最长的一条边作为长
    length = max(edges)
    longest_edge_index = edges.index(length)

    # 确定长边的两个顶点
    if longest_edge_index == 0:
        p1, p2 = points[0], points[1]
    elif longest_edge_index == 1:
        p1, p2 = points[1], points[2]
    elif longest_edge_index == 2:
        p1, p2 = points[2], points[3]
    else:
        p1, p2 = points[3], points[0]

    # 计算宽（与长边垂直的方向上的最大距离）
    width = max(
        perpendicular_distance(points[2], p1, p2),
        perpendicular_distance(points[3], p1, p2)
    )
    
    # print(f"Length: {length}")
    # print(f"Width: {width}")
    return length,width